from PyQt5.QtWidgets import QFileDialog, QInputDialog, QMessageBox, QPushButton, QAction
from PyQt5.uic import loadUi
import pyqtgraph as pg

import numpy as np
import pandas as pd
import cv2 as cv

from main import Homography_8, Homography_6, Homography_4

class MainWindow:

    def __init__(self) -> None:
        from PyQt5.QtWidgets import QWidget
        self.ui: QWidget = loadUi('main.ui')
        action_random: QAction = self.ui.action_random
        action_load: QAction = self.ui.action_load
        action_auto: QAction = self.ui.action_auto
        action_4: QAction = self.ui.action_4
        action_6: QAction = self.ui.action_6
        action_8: QAction = self.ui.action_8
        action_transform: QAction = self.ui.action_transform
        action_introduction: QAction = self.ui.action_introduction
        Button_reg: QPushButton = self.ui.Button_reg
        Button_ref: QPushButton = self.ui.Button_ref

        action_random.triggered.connect(self.set_random_kpts)
        action_load.triggered.connect(self.set_loaded_kpts)
        action_auto.triggered.connect(self.set_auto_kpts)
        action_4.triggered.connect(self.registration_4)
        action_6.triggered.connect(self.registration_6)
        action_8.triggered.connect(self.registration_8)
        action_transform.triggered.connect(self.transform)
        action_introduction.triggered.connect(self.introduction)
        Button_reg.clicked.connect(self.open_image_reg)
        Button_ref.clicked.connect(self.open_image_ref)

        self.num = 0  # 特征点的数量
        # 设置x轴和y轴
        graphicsView: pg.GraphicsView = self.ui.graphicsView_reg
        graphicsView.setBackground(pg.mkColor((0, 0, 0)))
        self.viewBox_reg = pg.ViewBox(invertY=True, lockAspect=True)  # y轴翻转，同时锁定两个方向的缩放比
        graphicsView.setCentralItem(self.viewBox_reg)
        inf_y = pg.InfiniteLine(movable=False, angle=90, label='y ↓', pen=pg.mkColor((255, 0, 0)),
                            labelOpts={'position': 0.05, 'color': (0, 0, 0), 'fill': (255, 255, 255), 'movable': True})
        inf_x = pg.InfiniteLine(movable=False, angle=0, label='x →', pen=pg.mkColor((0, 255, 0)),
                            labelOpts={'position': 0.05, 'color': (0, 0, 0), 'fill': (255, 255, 255), 'movable': True})
        self.viewBox_reg.addItem(inf_y)
        self.viewBox_reg.addItem(inf_x)

        # 设置x轴和y轴
        graphicsView: pg.GraphicsView = self.ui.graphicsView_ref
        graphicsView.setBackground(pg.mkColor((0, 0, 0)))
        self.viewBox_ref = pg.ViewBox(invertY=True, lockAspect=True)  # y轴翻转，同时锁定两个方向的缩放比
        graphicsView.setCentralItem(self.viewBox_ref)
        inf_y = pg.InfiniteLine(movable=False, angle=90, label='y ↓', pen=pg.mkColor((255, 0, 0)),
                            labelOpts={'position': 0.05, 'color': (0, 0, 0), 'fill': (255, 255, 255), 'movable': True})
        inf_x = pg.InfiniteLine(movable=False, angle=0, label='x →', pen=pg.mkColor((0, 255, 0)),
                            labelOpts={'position': 0.05, 'color': (0, 0, 0), 'fill': (255, 255, 255), 'movable': True})
        self.viewBox_ref.addItem(inf_y)
        self.viewBox_ref.addItem(inf_x)

        self.imageItem_reg = pg.ImageItem(axisOrder='row-major')  # 显示图像的控件，默认'col-major'
        self.imageItem_ref = pg.ImageItem(axisOrder='row-major')
        self.viewBox_reg.addItem(self.imageItem_reg)
        self.viewBox_ref.addItem(self.imageItem_ref)
        self.target_reg_list = []  # 记录特征点的位置
        self.target_ref_list = []
        self.H: np.ndarray = None  # 保留配准的单应性矩阵

    def open_image_reg(self):
        """
        打开配准影像
        """
        from PyQt5.QtWidgets import QLineEdit
        file_path, filiter = QFileDialog.getOpenFileName(None, filter='*jpg;*png;*tif;;*jpg;;*png;;*tif')
        if not file_path:
            return
        lineEdit: QLineEdit = self.ui.lineEdit_reg
        lineEdit.setText(file_path)
        image = cv.imdecode(np.fromfile(file_path),  cv.IMREAD_ANYCOLOR)  # 使用OpenCV读取中文路径的图片
        image = cv.cvtColor(image, cv.COLOR_BGR2RGB)
        self.imageItem_reg.setImage(image, 256)

    def open_image_ref(self):
        """
        打开参考影像
        """
        from PyQt5.QtWidgets import QLineEdit
        file_path, filiter = QFileDialog.getOpenFileName(None, filter='*jpg;*png;*tif;;*jpg;;*png;;*tif')
        if not file_path:
            return
        lineEdit: QLineEdit = self.ui.lineEdit_ref
        lineEdit.setText(file_path)
        image = cv.imdecode(np.fromfile(file_path),  cv.IMREAD_ANYCOLOR)  # 使用OpenCV读取中文路径的图片
        image = cv.cvtColor(image, cv.COLOR_BGR2RGB)
        self.imageItem_ref.setImage(image, 256)

    def set_random_kpts(self):
        """
        随机生成特征点
        """
        num, info = QInputDialog.getInt(None, '设置特征点', '输入数量', self.num, 0, 600)
        if not info:
            return
        self.num = num
        # 删除所有特征点
        for item in self.target_reg_list:
            self.viewBox_reg.removeItem(item)
        for item in self.target_ref_list:
            self.viewBox_ref.removeItem(item)
        self.target_reg_list = []
        self.target_ref_list = []
        pts = np.random.uniform(0, 500, (self.num, 2))
        self.set_kpts(pts, pts)

    def set_loaded_kpts(self):
        """
        导入特征点
        """
        file_path, filter = QFileDialog.getOpenFileName(parent=None, filter='*xlsx')
        if not file_path:
            return
        information = pd.read_excel(file_path, index_col=0)
        self.num = information.shape[0]
        # 删除所有特征点
        for item in self.target_reg_list:
            self.viewBox_reg.removeItem(item)
        for item in self.target_ref_list:
            self.viewBox_ref.removeItem(item)
        self.target_reg_list = []
        self.target_ref_list = []
        try:
            pts_reg = information[['reg_x', 'reg_y']].values
            pts_ref = information[['ref_x', 'ref_y']].values
        except KeyError:
            return
        self.set_kpts(pts_reg, pts_ref)

    def set_auto_kpts(self):
        """
        自动生成匹配的特征点
        """
        num, info = QInputDialog.getInt(None, '提取特征点', '输入数量', value=500, min=0, max=600)
        if not info:
            return
        # 删除所有特征点
        for item in self.target_reg_list:
            self.viewBox_reg.removeItem(item)
        for item in self.target_ref_list:
            self.viewBox_ref.removeItem(item)
        self.target_reg_list = []
        self.target_ref_list = []
        
        orb = cv.ORB_create(num)
        bf = cv.BFMatcher(cv.NORM_HAMMING, crossCheck=False)
        # 读取输入图像
        image_reg = self.imageItem_reg.image
        image_ref = self.imageItem_ref.image
        image_reg = cv.cvtColor(image_reg, cv.COLOR_RGB2GRAY) if image_reg.shape == 3 else image_reg  # rgb格式转灰度
        image_ref = cv.cvtColor(image_ref, cv.COLOR_RGB2GRAY) if image_ref.shape == 3 else image_ref
        # 检测关键点和计算描述符
        kps_reg, des_reg = orb.detectAndCompute(image_reg, None)
        kps_ref, des_ref = orb.detectAndCompute(image_ref, None)
        matches = bf.knnMatch(des_reg, des_ref, k=2)  # 匹配关键点
        # 粗匹配
        good_matches = []
        for m, n in matches:
            if m.distance < 0.75 * n.distance:
                good_matches.append([m])
        # 统计特征点坐标
        src_pts = np.float32([kps_reg[m[0].queryIdx].pt for m in good_matches]).reshape(-1, 2)
        dst_pts = np.float32([kps_ref[m[0].trainIdx].pt for m in good_matches]).reshape(-1, 2)
        try:
            H, _ = cv.findHomography(src_pts, dst_pts, cv.RANSAC, 5.0)  # RANSAC算法精匹配
            # 统计特征点
            pts_reg = []
            pts_ref = []
            for i in range(_.shape[0]):
                if _[i] == 1:
                    pts_reg.append([src_pts[i, 0], src_pts[i, 1]])
                    pts_ref.append([dst_pts[i, 0], dst_pts[i, 1]])
            pts_reg = np.array(pts_reg)
            pts_ref = np.array(pts_ref)
        except cv.error:
            QMessageBox.warning(None, '提示', '粗匹配特征点数量不足4个', QMessageBox.Ok)
        except TypeError:
            QMessageBox.warning(None, '提示', '精匹配特征点数量不足4个', QMessageBox.Ok)
        finally:
            pts_reg = src_pts
            pts_ref = dst_pts
        self.num = pts_reg.shape[0]
        self.set_kpts(pts_reg, pts_ref)

    def set_kpts(self, pts_reg: np.ndarray, pts_ref: np.ndarray):
        """
        设置特征点
        :param pts_reg: 配准点坐标
        :param pts_ref: 参考点坐标
        """
        from pyqtgraph.Qt import QtCore
        h = np.random.uniform(0, 1000, self.num)  # 生成关键点的颜色hsv色域
        h /= 1000
        s = 1
        v = 1
        for i in range(self.num):
            targetItem_reg = pg.TargetItem(
                pos=(pts_reg[i, 0], pts_reg[i, 1]),
                size=20,
                pen=pg.hsvColor(h[i], s, v),
                label=str(i) + "\nx:{0: .2f}" + "\ny:{1: .2f}",
                labelOpts={
                    "offset": QtCore.QPoint(15, 15),
                    "color": pg.hsvColor(h[i], s, v)
                    }
            )
            targetItem_ref = pg.TargetItem(
                pos=(pts_ref[i, 0], pts_ref[i, 1]),
                size=20,
                pen=pg.hsvColor(h[i], s, v),
                label=str(i) + "\nx:{0: .2f}" + "\ny:{1: .2f}",
                labelOpts={
                    "offset": QtCore.QPoint(15, 15),
                    "color": pg.hsvColor(h[i], s, v)
                    }
            )
            self.target_reg_list.append(targetItem_reg)
            self.target_ref_list.append(targetItem_ref)
            self.viewBox_reg.addItem(targetItem_reg)
            self.viewBox_ref.addItem(targetItem_ref)

    def get_pts(self):
        """
        按顺序统计坐标
        :return: 
            pts_reg: 返回配准坐标
            pts_ref: 返回参考坐标
        """
        pts = []
        for item in self.target_reg_list:
            item: pg.TargetItem
            x, y = item.pos()
            pts.append([x, y])
        pts_reg = np.array(pts)
        pts = []
        for item in self.target_ref_list:
            item: pg.TargetItem
            x, y = item.pos()
            pts.append([x, y])
        pts_ref = np.array(pts)
        return pts_reg, pts_ref

    def registration_4(self):
        """
        自由度4的配准
        """
        pts_reg, pts_ref = self.get_pts()
        try:
            result = Homography_4(pts_reg, pts_ref)
        except Homography_4.error:
            QMessageBox.warning(None, '提示', '最少需要2个特征点', QMessageBox.Ok)
            return
        result.get_homography()
        self.H = result.H
        dw = Information_Dialog(result)
        dw.ui.exec_()

    def registration_6(self):
        """
        自由度6的配准
        """
        pts_reg, pts_ref = self.get_pts()
        try:
            result = Homography_6(pts_reg, pts_ref)
        except Homography_6.error:
            QMessageBox.warning(None, '提示', '最少需要3个特征点', QMessageBox.Ok)
            return
        result.get_homography()
        self.H = result.H
        dw = Information_Dialog(result)
        dw.ui.exec_()
    
    def registration_8(self):
        """
        自由度8的配准
        """
        pts_reg, pts_ref = self.get_pts()
        try:
            result = Homography_8(pts_reg, pts_ref)
        except Homography_8.error:
            QMessageBox.warning(None, '提示', '最少需要4个特征点', QMessageBox.Ok)
            return
        result.get_homography()
        self.H = result.H
        dw = Information_Dialog(result)
        dw.ui.exec_()
    
    def transform(self):
        """
        变换配准影像到参考影像的坐标系，并保存
        """
        image_reg = self.imageItem_reg.image
        image_ref = self.imageItem_ref.image
        if image_reg is None:
            QMessageBox.warning(None, '提示', '未打开配准图像', QMessageBox.Ok)
            return
        if image_ref is None:
            QMessageBox.warning(None, '提示', '未打开参考图像', QMessageBox.Ok)
            return
        if self.H is None:
            QMessageBox.warning(None, '提示', '没有选择自由度配准', QMessageBox.Ok)
            return
        image_trans = cv.warpPerspective(image_reg, self.H, dsize=(image_ref.shape[1], image_ref.shape[0]))
        file_path, filiter = QFileDialog.getSaveFileName(None, directory='', filter='*png;*jpg')
        if not file_path:
            return
        image_trans = cv.cvtColor(image_trans, cv.COLOR_RGB2BGR)
        retval, img_encoded = cv.imencode(file_path,  image_trans)
        img_encoded.tofile(file_path)  # 可保存中文路径

    def introduction(self):
        """
        介绍
        """
        text = ''
        text += '      第一步，设置特征点：包括随机设置、导入设置和自动设置。\n'
        text += '随机设置会随机放置一定数量的特征点，相同匹配的特征点颜色和编号相同；\n'
        text += '导入设置会导入之前保存的信息；\n'
        text += '自动设置需要打开配准图像和参考图像，会用ORB算子提取一定数量的特征点并匹配。\n\n'
        text += '      第二步，计算配准量：包括3种自由度，相应的自由度必须有一定的特征点数量，最后查看特征点的配准误差等结果。\n'
        text += '自由度4的配准量，包括旋转量（角度制）、缩放量、平移量；\n'
        text += '自由度6和8的配准量，分别包括单应性矩阵的前6和8个元素。\n\n'
        text += '      第三步，配准图像：将配准图像通过第二步计算的单应性矩阵变换到参考图像的坐标系，分辨率和参考图像相同。'
        QMessageBox.about(None, '介绍', text)

class Information_Dialog:
    def __init__(self, registration: Homography_8) -> None:
        """
        解析配准误差
        :param registeration: 配准对象
        """
        from PyQt5.QtWidgets import QDialog, QTextEdit, QLabel, QTableWidget, QTableWidgetItem
        self.ui: QDialog = loadUi('dialog.ui')
        self.ui.closeEvent = self.message
        pushButton_save: QPushButton = self.ui.pushButton_save
        textEdit: QTextEdit = self.ui.textEdit
        labelWidget: QLabel = self.ui.label
        tableWidget: QTableWidget = self.ui.tableWidget
        pushButton_save.clicked.connect(self.save_information)

        H = registration.H
        
        # 显示配准结果
        text = ''
        if isinstance(registration, Homography_4):
            scale = registration.scale
            angle = registration.angle
            shift_x = registration.shift_x
            shift_y = registration.shift_y
            text += 'scale:' + '\t' + str(scale) + '\n'
            text += 'angle:' + '\t' + str(angle) + '\n'
            text += 'shift_x:' + '\t' + str(shift_x) + '\n'
            text += 'shift_y:' + '\t' + str(shift_y) + '\n'
            labelWidget.setText('自由度4的配准结果：')
        elif isinstance(registration, Homography_6):
            labelWidget.setText('自由度6的配准结果：')
        else:
            labelWidget.setText('自由度8的配准结果：')
        text += 'H:' + '\n' + str(H)
        text = text.replace('[', ' ')
        text = text.replace(']', ' ')
        textEdit.setText(text)

        # 填充到表格控件中
        self.information: pd.DataFrame = registration.get_error()  # 读取信息
        tableWidget.setRowCount(self.information.shape[0])  # 可以研究排序功能
        for i in range(self.information.shape[0]):
            columns = ['reg_x', 'reg_y', 'trans_x', 'trans_y', 'ref_x', 'ref_y', 'error_x', 'error_y']
            for j, col in enumerate(columns):
                num = self.information.loc[i, col]
                item = QTableWidgetItem(f'{num: .5f}', type=0)
                tableWidget.setItem(i, j, item)
            
    def message(self, event):
        """
        窗口关闭操作
        :parms event: 
        """
        print(event)

    def save_information(self):
        """
        保存配准误差到xlsx文件
        """
        file_path, filiter = QFileDialog.getSaveFileName(None, caption='保存特征点的配准结果', directory='.xlsx', filter='*xlsx')
        if not file_path:
            return
        self.information.to_excel(file_path)

if __name__ == '__main__':
    from PyQt5.QtWidgets import QApplication
    app = QApplication([])
    mw = MainWindow()
    mw.ui.show()
    app.exec()
